package com.gitee.zyflzz.tool_set.bloomfilter;

import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicLongArray;
import java.util.function.ToIntFunction;
import java.util.stream.IntStream;

/**
 * 线程安全的布隆过滤器实现
 *
 * @param <T> 存储元素类型
 * @author Zhou Yifan
 */
public class BloomFilter<T> {

    /**
     * 使用 long 数组存储位信息 (每个 long 存储 64 位)
     */
    private final AtomicLongArray bits;
    /**
     * 总位数
     */
    private final int sizeInBits;
    /**
     * 哈希函数数量
     */
    private final int numHashFunctions;
    /**
     * 线程池
     */
    private final ExecutorService threadPool;

    /**
     * 当前已存储的数据量
     */
    private int size;

    /**
     * 哈希函数参数 (使用不同种子生成不同哈希)
     */
    private final ToIntFunction<T>[] hashFunctions;

    /**
     * 默认布隆过滤器容量大小 1MB
     */
    private static final int DEFAULT_SIZE_BYTES = 1 << 20;
    /**
     * 默认哈希函数数量 7
     */
    private static final int DEFAULT_HASH_NUM = 7;
    /**
     * 默认线程数量 2
     */
    private static final int DEFAULT_THREADS = 2;

    /**
     * 构造函数（使用默认参数）
     */
    public BloomFilter() {
        this(DEFAULT_SIZE_BYTES, DEFAULT_HASH_NUM, DEFAULT_THREADS, null);
    }

    /**
     * 完整参数构造函数
     *
     * @param sizeBytes        存储空间大小 (字节)
     * @param numHashFunctions 哈希函数数量
     * @param threadPoolSize   线程池大小
     * @param executorService  线程池, 若设置为 null, 则使用 <code>Executors.newFixedThreadPool(threadPoolSize)</code> 创建默认的线程池
     */
    @SuppressWarnings("unchecked")
    public BloomFilter(int sizeBytes, int numHashFunctions, int threadPoolSize, ExecutorService executorService) {
        // 参数校验
        if (sizeBytes > 1 << 30) throw new IllegalArgumentException("最大支持1GB存储");
        if (numHashFunctions < 1) throw new IllegalArgumentException("至少需要1个哈希函数");

        this.size = 0;
        this.sizeInBits = sizeBytes * 8;
        this.numHashFunctions = numHashFunctions;
        this.bits = new AtomicLongArray((sizeInBits + 63) / 64); // 向上取整
        if (executorService == null) this.threadPool = Executors.newFixedThreadPool(threadPoolSize);
        else this.threadPool = executorService;

        // 初始化哈希函数（使用不同种子）
        this.hashFunctions = new ToIntFunction[numHashFunctions];
        for (int i = 0; i < numHashFunctions; i++) {
            final int seed = i;
            hashFunctions[i] = obj -> Math.abs((obj.hashCode() ^ seed) * 16777619);
        }
    }

    /**
     * 添加元素 (线程安全)
     */
    public void add(T item) {
        for (ToIntFunction<T> hash : hashFunctions) {
            int hashValue = hash.applyAsInt(item);
            setBit(hashValue % sizeInBits);
        }
        // TODO: 这种方式可能存储在问题
        size++;
    }

    /**
     * 检查元素是否存在
     */
    public boolean mightContain(T item) {
        for (ToIntFunction<T> hash : hashFunctions) {
            int hashValue = hash.applyAsInt(item);
            if (!getBit(hashValue % sizeInBits)) {
                return false;
            }
        }
        return true;
    }

    /**
     * 批量添加 (并行处理)
     */
    public void addAll(Iterable<T> items) {
        items.forEach(item -> threadPool.submit(() -> add(item)));
    }

    /**
     * 原子方式设置位 (CAS操作)
     */
    private void setBit(int bitIndex) {
        int longIndex = bitIndex >>> 6; // 等价于 /64
        long mask = 1L << (bitIndex & 0x3F); // 等价于 %64

        while (true) {
            long current = bits.get(longIndex);
            long updated = current | mask;
            if (bits.compareAndSet(longIndex, current, updated)) {
                break;
            }
        }
    }

    /**
     * 原子方式读取位
     */
    private boolean getBit(int bitIndex) {
        int longIndex = bitIndex >>> 6;
        long mask = 1L << (bitIndex & 0x3F);
        return (bits.get(longIndex) & mask) != 0;
    }

    /**
     * 关闭资源 (必须调用)
     */
    public void shutdown() {
        threadPool.shutdown();
    }

    /**
     * 获取已经存储了多少数据
     */
    public int getSize() {
        return size;
    }

    public int getNumHashFunctions() {
        return numHashFunctions;
    }

    /**
     * 获取容量大小
     */
    public int getSizeInBits() {
        return sizeInBits;
    }
}
